import gradio as gr
from ultralytics import YOLO
import PIL.Image as Image
import json
import datetime

# 定义模型路径
MODEL_PATHS = {
    "模型1（默认）": "/root/autodl-tmp/ultralytics-yolo11-main/runs/train/exp23/weights/best.pt",
    "模型2": "/root/autodl-tmp/ultralytics-yolo11-main/runs/train/exp15/weights/best.pt",
    "模型3": "/root/autodl-tmp/ultralytics-yolo11-main/runs/train/exp17/weights/best.pt",
}

# 病害信息数据库
DISEASE_INFO = {
    "炭疽病": {
        "症状": "叶片出现褐色斑点，逐渐扩大并形成同心轮纹。",
        "防治方法": "及时清除病叶，喷洒多菌灵等杀菌剂。"
    },
    "黑斑病": {
        "症状": "果实表面出现黑色斑点，逐渐扩大并凹陷。",
        "防治方法": "加强果园通风，喷洒波尔多液等杀菌剂。"
    },
    # 其他病害信息
}

# 核桃品种信息数据库
WALNUT_VARIETIES = {
    "Bilecik": {
        "特点": "高产、抗病性强",
        "常见病害": ["炭疽病", "黑斑病"],
        "防治建议": "定期修剪，保持通风。"
    },
    "Chandler": {
        "特点": "果实大、壳薄",
        "常见病害": ["炭疽病", "黑斑病"],
        "防治建议": "加强果园管理，定期喷洒杀菌剂。"
    },
    "Fernor": {
        "特点": "适应性强、耐寒",
        "常见病害": ["炭疽病"],
        "防治建议": "及时清除病叶，保持果园清洁。"
    },
    # 其他品种信息
}

# 加载默认模型
default_model = YOLO(MODEL_PATHS["模型1（默认）"])

# 加载品种识别模型
variety_model = YOLO("/path/to/walnut_variety_model.pt")

# 保存历史记录
def save_history(image_path, disease_info, severity, advice, variety_name):
    history_entry = {
        "timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "image_path": image_path,
        "variety": variety_name,
        "disease_info": disease_info,
        "severity": severity,
        "advice": advice
    }
    with open("history.json", "a") as f:
        f.write(json.dumps(history_entry, ensure_ascii=False) + "\n")

# 查看历史记录
def view_history():
    try:
        with open("history.json", "r") as f:
            history = [json.loads(line) for line in f.readlines()]
        return history
    except FileNotFoundError:
        return []

# 病害严重程度评估
def evaluate_severity(results):
    severity_levels = ["轻微", "中等", "严重"]
    for r in results:
        total_area = r.orig_shape[0] * r.orig_shape[1]  # 图像总面积
        disease_area = sum(w * h for w, h in r.boxes.xywh.tolist())  # 病害区域总面积
        severity_ratio = disease_area / total_area
        
        if severity_ratio < 0.1:
            return severity_levels[0], "建议定期观察，无需立即处理。"
        elif severity_ratio < 0.3:
            return severity_levels[1], "建议采取防治措施，如喷洒杀菌剂。"
        else:
            return severity_levels[2], "病害严重，建议立即处理并加强果园管理。"

# 预测函数
def predict_image(img, conf_threshold, iou_threshold, model_name):
    # 识别核桃品种
    variety_name = predict_variety(img)
    variety_info = WALNUT_VARIETIES.get(variety_name, {})
    
    # 根据用户选择的模型名称加载对应的病害识别模型
    if model_name in MODEL_PATHS:
        model = YOLO(MODEL_PATHS[model_name])
    else:
        model = default_model
    
    # 病害识别
    results = model.predict(
        source=img, 
        conf=conf_threshold, 
        iou=iou_threshold, 
        show_labels=True, 
        show_conf=True,
        imgsz=640,
    )
    
    for r in results:
        im_array = r.plot()
        im = Image.fromarray(im_array[..., ::-1])
        # 获取病害分类结果
        detected_diseases = set(r.boxes.cls.tolist())  # 获取检测到的病害类别
        disease_info = {}
        for disease_id in detected_diseases:
            disease_name = model.names[int(disease_id)]  # 获取病害名称
            if disease_name in DISEASE_INFO:
                disease_info[disease_name] = DISEASE_INFO[disease_name]
        
        # 评估病害严重程度
        severity, advice = evaluate_severity(results)
        
        # 保存历史记录
        save_history("uploaded_image.jpg", disease_info, severity, advice, variety_name)
    
    return im, disease_info, severity, advice, variety_info

# 主题语
theme_description = """
<p style="text-align: center;">守护核桃健康，助力农业丰收！</p>
<p style="text-align: center;">上传核桃图片，AI智能识别病害，助您及时防治！</p>
"""

# Gradio界面
iface = gr.Interface(
    fn=predict_image,
    inputs=[
        gr.Image(type="pil", label="上传图片", height=280, width=960),
        gr.Slider(minimum=0, maximum=1, value=0.25, label="置信度阈值"),
        gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU 阈值"),
        gr.Dropdown(choices=list(MODEL_PATHS.keys()), label="选择模型", value="模型1（默认）"),
    ],
    outputs=[
        gr.Image(type="pil", label="识别结果", height=280, width=960),
        gr.JSON(label="病害信息"),
        gr.Textbox(label="病害严重程度"),
        gr.Textbox(label="防治建议"),
        gr.JSON(label="品种信息")
    ],
    title="“核”护有家—核桃病害AI识别系统",
    description=theme_description,
    examples=[
        ["/root/autodl-tmp/Dataset/train/chandler/IMG_3088.JPG", 0.25, 0.45, "模型1（默认）"],
        ["/root/autodl-tmp/Dataset/train/chandler/IMG_3095.JPG", 0.25, 0.45, "模型1（默认）"],
    ],
    theme="soft",
    allow_flagging="never",
)

# 历史记录界面
history_iface = gr.Interface(
    fn=view_history,
    inputs=[],
    outputs=gr.JSON(label="历史记录"),
    title="历史记录",
    allow_flagging="never"
)

# 启动界面
if __name__ == '__main__':
    gr.TabbedInterface([iface, history_iface], ["病害识别", "历史记录"]).launch()